import tempfile
from collections import defaultdict
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Tuple, Optional
import flatbuffers

import cv2
import numpy as np
from mcap.reader import make_reader
from mcap.well_known import MessageEncoding
from scipy.spatial.transform import Rotation as R


PROJECT_ROOT = Path(__file__).parent.parent.parent
DEFAULT_TASK_INFO_PATH = PROJECT_ROOT / "task_info.json"

EEF_OFFSET_X = np.array(
    [
        [1.0, 0.0, 0.0, 0.1],
        [0.0, 1.0, 0.0, 0.0],
        [0.0, 0.0, 1.0, 0.0],
        [0.0, 0.0, 0.0, 1.0],
    ],
    dtype=np.float32,
)


class FloatArray:
    __slots__ = ["_tab"]

    @classmethod
    def GetRootAs(cls, buf, offset=0):
        n = flatbuffers.encode.Get(flatbuffers.packer.uoffset, buf, offset)
        x = FloatArray()
        x.Init(buf, n + offset)
        return x

    @classmethod
    def GetRootAsFloatArray(cls, buf, offset=0):
        """This method is deprecated. Please switch to GetRootAs."""
        return cls.GetRootAs(buf, offset)

    def Init(self, buf, pos):
        self._tab = flatbuffers.table.Table(buf, pos)

    def Values(self, j):
        o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
        if o != 0:
            a = self._tab.Vector(o)
            return self._tab.Get(
                flatbuffers.number_types.Float32Flags,
                a + flatbuffers.number_types.UOffsetTFlags.py_type(j * 4),
            )
        return 0

    def ValuesAsNumpy(self):
        o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
        if o != 0:
            return self._tab.GetVectorAsNumpy(flatbuffers.number_types.Float32Flags, o)
        return 0

    def ValuesLength(self):
        o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
        if o != 0:
            return self._tab.VectorLen(o)
        return 0

    def ValuesIsNone(self):
        o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
        return o == 0


@dataclass(frozen=True)
class SyncedMCAPData:
    """
    用于存储从 MCAP 文件同步后的数据帧。
    所有列表/数组的第一个维度（长度）都是同步后的有效帧数 (T)。
    """

    timestamps: List[int]  # [T]
    images: Dict[str, List[np.ndarray]]  # {"cam_key": [T, H, W, C]}
    qpos_obs: np.ndarray  # [T, D_qpos_obs]
    qpos_act: np.ndarray  # [T, D_qpos_act]
    gpos_obs: np.ndarray  # [T, D_gpos_obs]
    gpos_act: np.ndarray  # [T, D_gpos_act]


class MCAPBaseReader:
    """
    从 MCAP 文件读取和同步机器人数据（图像、关节状态、姿态）的基类。

    它负责处理所有底层的 MCAP 解析、视频解码和时间戳同步。
    子类可以继承此类并覆盖 `read_and_sync_data()` 返回的数据，
    以实现特定的文件格式转换（HDF5, LeRobot, 视频等）。
    """

    # --- 共享常量 ---
    CAMERA_MAPPING = {
        "cam_left": "left_camera",
        "cam_head": "env_camera",
        "cam_right": "right_camera",
    }
    ATTACHMENT_TO_CAM = {
        "/left_camera/color/image_raw": "left_camera",
        "/env_camera/color/image_raw": "env_camera",
        "/right_camera/color/image_raw": "right_camera",
    }
    DEFAULT_JOINT_DOF = 7

    def __init__(self, mcap_file: Path):
        self.mcap_file = mcap_file
        self.camera_names = list(self.CAMERA_MAPPING.keys())  # hdf5/lerobot names

        # --- 共享的 Topic 字段定义 ---

        # 关节角度 (qpos) 数据
        self.obs_fields = [
            "/left/follow/arm/joint_state/position",
            "/left/follow/eef/joint_state/position",
            "/right/follow/arm/joint_state/position",
            "/right/follow/eef/joint_state/position",
        ]
        self.act_fields = [
            "/left/lead/arm/joint_state/position",
            "/left/lead/eef/joint_state/position",
            "/right/lead/arm/joint_state/position",
            "/right/lead/eef/joint_state/position",
        ]

        # 全局位姿 (gpos) 数据
        self.obs_gpos_fields = [
            "/left/follow/arm/pose/position",
            "/left/follow/arm/pose/orientation",
            "/left/follow/eef/joint_state/position",  # eef 只有1维
            "/right/follow/arm/pose/position",
            "/right/follow/arm/pose/orientation",
            "/right/follow/eef/joint_state/position",  # eef 只有1维
        ]
        self.act_gpos_fields = [
            "/left/lead/arm/pose/position",
            "/left/lead/arm/pose/orientation",
            "/left/lead/eef/joint_state/position",
            "/right/lead/arm/pose/position",
            "/right/lead/arm/pose/orientation",
            "/right/lead/eef/joint_state/position",
        ]

    def read_and_sync_data(self) -> Optional[SyncedMCAPData]:
        """
        公共方法：执行完整的读取和同步流程。
        """
        try:
            with open(self.mcap_file, "rb") as f:
                reader = make_reader(f)
                joint_data_by_ts = self._read_joint_data(reader)
                image_data, cam_frame_counts = self._read_image_data(reader)

            if not joint_data_by_ts or not image_data:
                print(
                    f"[*] 警告: 文件 {self.mcap_file.name} 缺少有效的关节或图像数据。"
                )
                return None

            synced_data = self._prepare_synced_datasets(
                joint_data_by_ts, image_data, cam_frame_counts
            )

            if synced_data is None:
                print(f"[*] 警告: 文件 {self.mcap_file.name} 同步后无有效数据帧。")
                return None

            return synced_data

        except Exception as e:
            print(f"[!] 错误: 处理 MCAP 文件 {self.mcap_file.name} 时发生异常: {e}")
            return None

    # --- 核心私有方法 ---

    def _read_joint_data(self, reader) -> Dict[int, Dict[str, np.ndarray]]:
        """读取所有相关 topic 的关节数据，按时间戳组织"""
        timestamped_data = defaultdict(dict)
        # 读取所有已定义的 topics (qpos + gpos)
        all_topics_to_read = list(
            set(
                self.obs_fields
                + self.act_fields
                + self.obs_gpos_fields
                + self.act_gpos_fields
            )
        )

        for schema, channel, message in reader.iter_messages(topics=all_topics_to_read):
            if schema and schema.encoding == MessageEncoding.Flatbuffer:
                timestamped_data[message.log_time][channel.topic] = (
                    self._parse_float_array(message.data)
                )
        return timestamped_data

    @staticmethod
    def _parse_float_array(data: bytes) -> np.ndarray:
        float_array = FloatArray.GetRootAs(data, 0)
        return float_array.ValuesAsNumpy()

    def _read_image_data(
        self, reader
    ) -> Tuple[Dict[str, Dict[int, np.ndarray]], Dict[str, int]]:
        """读取所有摄像头的视频附件数据"""
        image_data_by_cam = defaultdict(dict)
        camera_frame_counts = {}
        for attachment in reader.iter_attachments():
            if attachment.name not in self.ATTACHMENT_TO_CAM:
                continue
            camera_key = self.ATTACHMENT_TO_CAM[attachment.name]
            frames = self._decode_video(attachment.data)
            if not frames:
                continue
            camera_frame_counts[camera_key] = len(frames)
            for i, frame in enumerate(frames):
                image_data_by_cam[camera_key][i] = frame
        return image_data_by_cam, camera_frame_counts

    @staticmethod
    def _decode_video(video_data: bytes) -> List[np.ndarray]:
        """将 bytes 格式的视频附件解码为 RGB 帧列表"""
        frames = []
        # 使用 delete=True 更安全，自动清理临时文件
        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=True) as temp_f:
            temp_f.write(video_data)
            temp_f.flush()
            cap = cv2.VideoCapture(temp_f.name)
            while cap.isOpened():
                ret, frame = cap.read()
                if not ret:
                    break
                frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
            cap.release()
        return frames

    def _apply_eef_correction(
        self, gpos_array: np.ndarray, X: np.ndarray
    ) -> np.ndarray:
        """
        对 act_gpos (主臂) 的【左臂和右臂】末端姿态同时做齐次变换修正。
        gpos_array: shape [N, D]，其中 D 包含左右臂的 (pos+rpy+eef)
        X: 4x4 齐次矩阵（从主臂TCP到从臂TCP的固定变换）
        """
        corrected = gpos_array.astype(np.float32)
        num_frames = gpos_array.shape[0]

        # 定义左右臂的索引切片
        # GPOS 阵列布局:
        # [L_arm_pos(3), L_arm_rpy(3), L_eef(1), R_arm_pos(3), R_arm_rpy(3), R_eef(1)]
        arm_indices = [
            {"pos": slice(0, 3), "rpy": slice(3, 6)},  # Left arm
            {"pos": slice(7, 10), "rpy": slice(10, 13)},  # Right arm
        ]

        T = np.eye(4, dtype=np.float32)  # 预先分配齐次矩阵，在循环中重用

        for i in range(num_frames):
            # 对定义的每个臂（左、右）都应用变换
            for arm in arm_indices:
                pos_slice = arm["pos"]
                rpy_slice = arm["rpy"]

                pos = corrected[i, pos_slice]
                rpy = corrected[i, rpy_slice]

                # RPY → 齐次矩阵
                T[:3, :3] = R.from_euler("xyz", rpy).as_matrix()
                T[:3, 3] = pos

                # 应用偏移变换
                T_new = T @ X

                # 提取回 RPY + pos
                corrected[i, pos_slice] = T_new[:3, 3]
                corrected[i, rpy_slice] = (
                    R.from_matrix(T_new[:3, :3]).as_euler("xyz").astype(np.float32)
                )

        return corrected

    def _prepare_synced_datasets(
        self, joint_data, image_data, cam_frame_counts
    ) -> Optional[SyncedMCAPData]:
        """
        同步关节数据和图像数据。
        返回一个包含所有同步后数据的 dataclass。
        """
        sorted_timestamps = sorted(joint_data.keys())
        min_image_frames = min(cam_frame_counts.values()) if cam_frame_counts else 0
        effective_len = min(len(sorted_timestamps), min_image_frames)

        if effective_len == 0:
            return None

        final_timestamps = sorted_timestamps[:effective_len]

        # images_dict: {"left_camera": [frame1, ...], ...}
        images_dict = {cam_key: [] for cam_key in self.ATTACHMENT_TO_CAM.values()}

        obs_qpos_list, act_qpos_list = [], []
        obs_gpos_list, act_gpos_list = [], []

        # 用于在帧丢失时填充上一帧
        last_frame_cache = {cam_key: None for cam_key in images_dict}
        default_frame = np.zeros((480, 640, 3), dtype=np.uint8)

        for frame_idx, ts in enumerate(final_timestamps):
            joint_values_at_ts = joint_data.get(ts, {})

            # 1. 处理 QPOS 数据
            obs_qpos_list.append(
                self._get_concatenated_values(joint_values_at_ts, self.obs_fields)
            )
            act_qpos_list.append(
                self._get_concatenated_values(joint_values_at_ts, self.act_fields)
            )

            # 2. 处理 GPOS 数据
            obs_gpos_list.append(
                self._get_concatenated_gpos(joint_values_at_ts, self.obs_gpos_fields)
            )
            act_gpos_list.append(
                self._get_concatenated_gpos(joint_values_at_ts, self.act_gpos_fields)
            )

            # 3. 处理图像数据
            for mcap_cam_key in images_dict.keys():
                frame = image_data.get(mcap_cam_key, {}).get(frame_idx)
                if frame is None:
                    # 优先使用上一帧，如果上一帧也没有，使用黑屏
                    frame = (
                        last_frame_cache[mcap_cam_key]
                        if last_frame_cache[mcap_cam_key] is not None
                        else default_frame
                    )

                images_dict[mcap_cam_key].append(frame)
                last_frame_cache[mcap_cam_key] = frame

        # 末端姿态补偿
        try:
            act_gpos_list_np = np.array(act_gpos_list, dtype=np.float32)
            act_gpos_list_np = self._apply_eef_correction(
                act_gpos_list_np, EEF_OFFSET_X
            )
        except Exception as e:
            print(f"[!] 末端姿态修正失败: {e}")
            act_gpos_list_np = np.array(act_gpos_list, dtype=np.float32)

        return SyncedMCAPData(
            timestamps=final_timestamps,
            images=images_dict,
            qpos_obs=np.array(obs_qpos_list, dtype=np.float32),
            qpos_act=np.array(act_qpos_list, dtype=np.float32),
            gpos_obs=np.array(obs_gpos_list, dtype=np.float32),
            gpos_act=act_gpos_list_np,
        )

    # --- 辅助方法 ---

    def _get_concatenated_values(
        self, data_dict: Dict, fields: List[str]
    ) -> np.ndarray:
        """串联 qpos 数据"""
        default_val = np.zeros(self.DEFAULT_JOINT_DOF, dtype=np.float32)
        arrays_to_concat = [data_dict.get(field, default_val) for field in fields]
        return np.concatenate(arrays_to_concat)

    @staticmethod
    def _quat_to_rpy(quat_array: np.ndarray) -> np.ndarray:
        """将四元数 [x, y, z, w] 转换为 rpy [roll, pitch, yaw]"""
        if quat_array.shape[-1] == 4:
            r = R.from_quat(quat_array)
            return r.as_euler("xyz", degrees=False)
        else:
            # 如果数据不是四元数（例如，已经是rpy或无效数据），返回0
            return np.zeros(3, dtype=np.float32)

    def _get_concatenated_gpos(self, data_dict: Dict, fields: List[str]) -> np.ndarray:
        """串联 gpos 数据，并处理四元数到RPY的转换"""
        processed_arrays = []
        for field in fields:
            raw_array = data_dict.get(field)

            # --- 处理数据缺失 ---
            if raw_array is None:
                if "orientation" in field:
                    # 姿态 RPY (roll, pitch, yaw)
                    processed_arrays.append(np.zeros(3, dtype=np.float32))
                elif "position" in field and "eef" not in field:
                    # 位置 XYZ
                    processed_arrays.append(np.zeros(3, dtype=np.float32))
                elif "eef" in field and "position" in field:
                    # 末端夹爪（通常是1维）
                    processed_arrays.append(np.zeros(1, dtype=np.float32))
                else:
                    # 兜底，防止意外
                    processed_arrays.append(np.array([], dtype=np.float32))
                continue

            # --- 处理有效数据 ---
            if "orientation" in field:
                processed_arrays.append(self._quat_to_rpy(raw_array))
            else:
                processed_arrays.append(raw_array)

        return np.concatenate(processed_arrays)
